module Hexena.MZHeader
  ( MZHeader
  , checkFile
  , commonDefaultHeaderWord8
  , commonDefaultHeader
  , commonDefaultMZ
  , makeMZHD
  , loadMZHD
  , mzDistance
  , mzDistanceN
  ) where

import Prelude hiding (lookup)
import Data.Binary.Get
import Data.Binary.Put
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as L
import Data.IFF
import Data.Map (Map, insert, empty, lookup)
import Data.Maybe
import Data.Word

type MZHeader = (Map String Word16, Word32)

-- | Keys for the MZ Header.
keys :: [String]
keys = [ "magic"
       , "cplp", "cp"
       , "crlc"
       , "cparhdr"
       , "minalloc", "maxalloc"
       , "ss"
       , "sp"
       , "csum"
       , "ip"
       , "cs"
       , "lfarlc"
       , "ovno"
       , "res_0", "res_1", "res_2", "res_3"
       , "oemid", "oeminfo"
       , "res2_0", "res2_1", "res2_2", "res2_3"
       , "res2_4", "res2_5", "res2_6", "res2_7"
       , "res2_8", "res2_9"
       ]

-- | Read the MZ header and return the lfanew field. This is currently
-- done as a map as it is easy and convenient. Not much is done with
-- the MZ header anyways.
checkFile :: Get MZHeader
checkFile = let getNext [] = return empty
                getNext (k:xs) = do
                  v <- getWord16le
                  r <- getNext xs
                  return $ insert k v r
            in do uss <- getNext keys
                  lfanew <- getWord32le
                  return (uss, lfanew)

makeMZHD :: Map String Word16 -> Word32 -> T
makeMZHD m lfanew = Cons (chunkIdFromString "MZHD") $ Chunk $ B.concat $ L.toChunks $ runPut putF
  where putF = do mapM_ (putWord16be . fromJust . flip lookup m) keys  
                  putWord32be lfanew

loadMZHD :: Chunk -> MZHeader -- (Map String Word16, Word32)
loadMZHD t = let c = L.fromChunks [content t]
                 getF = do xs <- get16 keys
                           ys <- getWord32be
                           return (xs, ys)
                 get16 [] = return empty
                 get16 (k:ks) = do v <- getWord16be
                                   r <- get16 ks
                                   return $ insert k v r
             in runGet getF c

-- This is a normal MZ header. On my Dell Vostro default Win7 install
-- this is the most common header (99% of the files on my
-- laptop). Therefore we calculate a "distance" function to this
-- normal header.

-- $ find /mnt/ -type f -iname *.exe -exec sh -c "xxd -c64 -l 64 '{}'" \;|grep "0000000: 4d5a"|wc
--    2954  100439  694190
-- $ find /mnt/ -type f -iname *.exe -exec sh -c "xxd -c64 -l 64 '{}'" \;|grep "0000000: 4d5a 9000 0300 0000 0400 0000 ffff 0000 b800 0000 0000 0000 4000 0000 0000 0000 0000 0000 0000 0000 0000 0000 0000 0000 0000 0000 0000 0000 0000 0000"|wc
--    2927   99521  687845
-- (/ (* 100.0 2927) 2954)

-- 0000000: 4d5a 9000 0300 0000 0400 0000 ffff 0000  MZ..............
-- 0000010: b800 0000 0000 0000 4000 0000 0000 0000  ........@.......
-- 0000020: 0000 0000 0000 0000 0000 0000 0000 0000  ................
-- 0000030: 0000 0000 0000 0000 0000 0000 e800 0000  ................


-- | Default MZ Header as a Word8 list.
commonDefaultHeaderWord8 :: [Word8]
commonDefaultHeaderWord8 = [ 0x4d, 0x5a, 0x90, 0x00, 0x03, 0x00, 0x00, 0x00, 
                             0x04, 0x00, 0x00, 0x00, 0xff, 0xff, 0x00, 0x00, 
                             0xb8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
                             0x40, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
                             0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
                             0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
                             0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
                             0x00, 0x00, 0x00, 0x00 ]

-- | Default MZ Header as a ByteString list.
commonDefaultHeader :: B.ByteString
commonDefaultHeader = B.pack commonDefaultHeaderWord8

commonDefaultMZ :: MZHeader
commonDefaultMZ = runGet checkFile $ L.pack $ commonDefaultHeaderWord8 ++ [0xFF, 0xFF, 0xFF, 0xFF]


-- | Calculate "euclidic distance" to a "normal" MZ header.
mzDistance :: [Word8] -> Double
mzDistance xs | length xs >= 60 = sqrt $ foldl f 0.0 $ zip commonDefaultHeaderWord8 xs
              | otherwise = error "MZ header has missing bytes in mzDistance"
  where f :: Double -> (Word8, Word8) -> Double
        f z (x, y) = z + (fromIntegral x - fromIntegral y) ^ (2::Int)

-- | Calculate normalised "euclidic distance" to a "normal" MZ header. Maximum is 1.
mzDistanceN :: [Word8] -> Double
mzDistanceN x = mzDistance x / 255 / sqrt 60
